#include "stdafx.h"

void Register() 
{
	int ret = RegisterTraceGuids(
        (WMIDPREQUEST)control,	// register control entry point
        0,	// context
        &guid1,
        1,	// corr count
        correlators,	// correlators
        NULL,			// mof image path
        NULL,			// mof resource name
        &registrationh
    );
}
void Unregister() 
{
    int ret = UnregisterTraceGuids(registrationh);
	return;
}

// Session management

void SessionInitialize(PEVENT_TRACE_PROPERTIES& loggerdata) 
{
    loggerdata_len = 
        sizeof (EVENT_TRACE_PROPERTIES) +
        lstrlenA(LoggerName) + 1 +
        lstrlenA(EtlLoggerFileName) + 1;
	loggerdata = (PEVENT_TRACE_PROPERTIES) malloc(loggerdata_len);  //todo check alloc
	memset(loggerdata,0,loggerdata_len);
	loggerdata->FlushTimer = 0;
    loggerdata->NumberOfBuffers = 50;
    loggerdata->BufferSize = 2048;
    loggerdata->MaximumFileSize = 200;
    loggerdata->Wnode.BufferSize = loggerdata_len;
    loggerdata->Wnode.Flags = WNODE_FLAG_TRACED_GUID; 
    loggerdata->Wnode.Guid = guid1;
    loggerdata->LoggerNameOffset = sizeof(EVENT_TRACE_PROPERTIES);
    loggerdata->LogFileNameOffset = loggerdata->LoggerNameOffset + lstrlenA(LoggerName) + 1;
	loggerdata->LogFileMode =   EVENT_TRACE_FILE_MODE_CIRCULAR | EVENT_TRACE_USE_PAGED_MEMORY;

    // Set logger name
    
    char* loggerName = (LPSTR)((char*)loggerdata + loggerdata->LoggerNameOffset); 
	StringCchCopyA(loggerName, lstrlenA(LoggerName) + 1, LoggerName);

    // Set log file name

    char* logfileName = (LPSTR)((char*)loggerdata + loggerdata->LogFileNameOffset); 
    StringCchCopyA(logfileName, lstrlenA(EtlLoggerFileName) + 1, EtlLoggerFileName);
}
void SessionStart(PEVENT_TRACE_PROPERTIES loggerdata) 
{
	 ULONG ret = StartTrace(&loggerh, LoggerName, loggerdata);
}
void SessionEnable() 
{
	ULONG ret = EnableTrace (TRUE, 0,0, &guid1, loggerh);
}
void SessionDisable() 
{
	ULONG ret = EnableTrace (FALSE, 0,0, &guid1, loggerh);
}
void ControlStop(PEVENT_TRACE_PROPERTIES loggerdata) 
{
	ULONG ret = ControlTrace(loggerh, LoggerName, loggerdata, EVENT_TRACE_CONTROL_STOP);
	if (ERROR_MORE_DATA == ret) printf ("ERROR_MORE_DATA\n");
}


// PROVIDER code

ULONG FireBasicFormat(const char* message) 
{
	GenericEvent.Header.Flags = WNODE_FLAG_TRACED_GUID;
	GenericEvent.Header.Guid  = CorrGuid;
    GenericEvent.Header.Class.Type = EVENT_TRACE_TYPE_INFO;

	StringCchCopyA(GenericEvent.CustomData, MAX_MESSAGE, message);
    GenericEvent.FormatCode = FORMAT_BASIC;
    GenericEvent.Header.Size  = sizeof(EVENT_TRACE_HEADER) + sizeof(GenericEvent.FormatCode) + (USHORT)lstrlenA(GenericEvent.CustomData) + 1;
    ULONG ret = TraceEvent(loggerh, (PEVENT_TRACE_HEADER) &GenericEvent);
    return ret;
}

// CONTROL callback

ULONG control( IN WMIDPREQUESTCODE requestCode, IN PVOID Context, IN OUT ULONG *InOutbufferSize, IN OUT PVOID buffer)
{
    switch (requestCode)
    {
        case WMI_ENABLE_EVENTS:
			{
            loggerh = GetTraceLoggerHandle( buffer );
            g_level = GetTraceEnableLevel(loggerh);
            g_flags = GetTraceEnableFlags(loggerh);
            enabled = TRUE;
			}
            break;
        case WMI_DISABLE_EVENTS:
            loggerh = GetTraceLoggerHandle( buffer );
            g_level = GetTraceEnableLevel(loggerh);
            enabled = FALSE;
            loggerh = 0;
            break;
        default:
			printf(" control: invalid request code: %d\n", (int)requestCode);
    }
    return 0;
}

// READING STUFF

ULONG ReadWholeTrace(const char* loggerName, const char* etlPath, const char* txtPath)
{
	EvmFile = (PEVENT_TRACE_LOGFILE) malloc(sizeof(EVENT_TRACE_LOGFILE));
    TRACEHANDLE HandleArray[MAXLOGFILES];
	RtlZeroMemory(EvmFile, sizeof(EVENT_TRACE_LOGFILE));
	EvmFile->LoggerName = (LPSTR)loggerName;
	EvmFile->Context = NULL;
	EvmFile->BufferCallback = BufferCallback;
	EvmFile->BuffersRead = 0;
	EvmFile->CurrentTime = 0;
	EvmFile->EventCallback = &DumpEvent;
	EvmFile->LogFileMode = EVENT_TRACE_FILE_MODE_SEQUENTIAL;
    EvmFile->LogFileName = (LPSTR)etlPath;

    TRACEHANDLE xxxxx = OpenTrace(EvmFile);
    if (xxxxx == 0) 
    {
        DWORD Status = GetLastError();
        printf("Error Opening Trace %d with status=%d\n", 0, Status);
    }
    HandleArray[0] = xxxxx;

    hDecodedFile = CreateFile(
                txtPath,
                GENERIC_WRITE,
                FILE_SHARE_READ | FILE_SHARE_WRITE,
                NULL,
                CREATE_ALWAYS,
                FILE_ATTRIBUTE_NORMAL,
                NULL
                );
    if (hDecodedFile == INVALID_HANDLE_VALUE)
    {
        return 666;
    }

    DWORD Status = ProcessTrace(
            HandleArray,
            1,  // handle count
            NULL,
            NULL
            );
    CloseHandle(hDecodedFile);

    if (Status != ERROR_SUCCESS) {
		_tprintf(_T("ProcessTrace failed: %dL (GetLastError=0x%x)\n"),
                Status, GetLastError());
    }

    for (int jiiii = 0; jiiii < sizeof(HandleArray) / sizeof(HandleArray[0]); jiiii++){
        DWORD Status = CloseTrace(HandleArray[jiiii]);
        if (Status != ERROR_SUCCESS) 
        {
            _tprintf(_T("Error Closing Trace %d with status=%d\n"), jiiii, Status);
        }
    }
    return NO_ERROR;
}

ULONG WINAPI BufferCallback(PEVENT_TRACE_LOGFILE pLog)
{
	printf(".");
    return (TRUE);
}

void WINAPI DumpEvent(PEVENT_TRACE pEvent)
{
    BOOL bUserMode = FALSE;
    ULONG PointerSize = sizeof(PVOID) * 8;
	//printf("#%d;", TotalEventCount);
    PEVENT_TRACE_HEADER pHeader;
//  PMOF_INFO pMofInfo;

    TotalEventCount++;

    if (pEvent == NULL) {
        _tprintf(_T("Warning: Null Event\n"));
        return;
    }

    pHeader = (PEVENT_TRACE_HEADER) &pEvent->Header;

    // Extrace log file information if the event is a log file header.
	if( IsEqualGUID(&(pEvent->Header.Guid), &EventTraceGuid)) {
		///printf ("EventTraceGUID\n");
	}

    if( IsEqualGUID(&(pEvent->Header.Guid), &EventTraceGuid) && pEvent->Header.Class.Type == EVENT_TRACE_TYPE_INFO ) 
    {
        PTRACE_LOGFILE_HEADER head = (PTRACE_LOGFILE_HEADER)pEvent->MofData;
        if( NULL != head ){
            if(head->TimerResolution > 0){
                TimerResolution = head->TimerResolution / 10000;
            }
        
            StartTime  = head->StartTime.QuadPart;
            EndTime    = head->EndTime.QuadPart;
            // If ProcessTrace() call was made on areal time logger or an trace file being 
            // logged, EndTime amy be 0.
            fNoEndTime = (EndTime == 0);

            TotalEventsLost += head->EventsLost;

            // We use global flags for private logger and pointer size.
            // This may cause an error when trace files are from different environments.
            bUserMode = (head->LogFileMode & EVENT_TRACE_PRIVATE_LOGGER_MODE);

            // Set pointer size
            PointerSize =  head->PointerSize * 8;
            if (PointerSize != 64){   
                PointerSize = 32; 
            }
        }
    }
    else if( IsEqualGUID(&(pEvent->Header.Guid), &CorrGuid) && pEvent->Header.Class.Type == EVENT_TRACE_TYPE_INFO ) 
    {
        char* data = (char*)pEvent->MofData + 4;
        int len = pEvent->Header.Size;
        char buffer[MAX_MESSAGE];
        // \n --> \r\n
        char* p = buffer;
        for (; *data && ((p - buffer) < (sizeof(buffer) - 3));)
        {
            if (*data == '\r')
            {
                *p++ = *data++; if (*data == 0) break;
                *p++ = *data++;
            }
            else if (*data == '\n')
            {
                *p++ = '\r';
                *p++ = *data++;
            }
            else
            {
                *p++ = *data++;
            }
        }
        *p = 0;
        // if no EOL, add one
        if ((p - buffer) > 0 && *(p-1) != '\n' && ((p - buffer) < (sizeof(buffer) - 3)))
        {
                *p++ = '\r';
                *p++ = '\n';
                *p++ = '\0';
        }
        len = lstrlenA(buffer);
        DWORD dw;
        BOOL ok = WriteFile(hDecodedFile, &buffer, len, &dw,0);
        if (!ok || len != dw)
        {
            printf ("* something wrong writing\n");
        }
    }
    else
    {
        printf ("* unknown event guid\n");
    }

    // if EndTime is missing from one of the files, keep updating to get the largest value.
    if (fNoEndTime && EndTime < (ULONGLONG) pHeader->TimeStamp.QuadPart) 
    {
        EndTime = pHeader->TimeStamp.QuadPart;
    }
}

